"""
PETSc-Only Extension Implementation.

This extension wraps pure PETSc solver implementations for use in GSM.
"""

import json
import os
import subprocess
import logging
from pathlib import Path
from typing import Dict, List, Optional, Any

from gcr_solver_manager.extensions.base_extension import SolverExtension
from gcr_solver_manager.extensions.base_builder import BuilderInterface
from gcr_solver_manager.extensions.base_runner import RunnerInterface
from gcr_solver_manager.extensions.base_validator import ValidatorInterface
from .builder import PETScOnlyBuilder
from .runner import PETScOnlyRunner
from .validator import PETScOnlyValidator

logger = logging.getLogger(__name__)


class PETScOnlyExtension(SolverExtension):
    """Extension for pure PETSc solver implementations."""
    
    def __init__(self, source_dir: Optional[Path] = None):
        """Initialize the PETSc-Only extension."""
        self.source_dir = source_dir or Path("../src/prototypes").resolve()
        self._builder = None
        self._runner = None
        self._validator = None
    
    @property
    def name(self) -> str:
        """Extension unique identifier."""
        return "petsc-only"
    
    @property
    def version(self) -> str:
        """Extension version string."""
        return "1.0.0"
    
    @property
    def supported_frameworks(self) -> List[str]:
        """List of supported framework names."""
        return ["petsc", "petsc-only"]
    
    @property
    def supported_solvers(self) -> List[str]:
        """List of supported solver types."""
        return ["gcr", "gmres", "cg", "bicg"]
    
    def get_builder(self) -> BuilderInterface:
        """Return builder instance for this extension."""
        if self._builder is None:
            self._builder = PETScOnlyBuilder(self.source_dir)
        return self._builder
    
    def get_runner(self) -> RunnerInterface:
        """Return runner instance for this extension."""
        if self._runner is None:
            self._runner = PETScOnlyRunner(self.source_dir)
        return self._runner
    
    def get_validator(self) -> ValidatorInterface:
        """Return validator instance for this extension."""
        if self._validator is None:
            self._validator = PETScOnlyValidator()
        return self._validator
    
    def get_config_schema(self) -> Dict[str, Any]:
        """Return JSON schema for configuration validation."""
        schema_path = Path(__file__).parent / "config_schema.json"
        if schema_path.exists():
            with open(schema_path, 'r') as f:
                return json.load(f)
        return {}
    
    def can_handle_config(self, config_path: Path) -> bool:
        """Check if extension can handle given configuration."""
        # Check if this is a PETSc-specific configuration
        
        if not config_path.exists():
            return False
            
        # Check if it's in the prototypes directory (but not hybrid or kokkos)
        try:
            if 'prototypes' in str(config_path):
                relative_path = config_path.resolve().relative_to(self.source_dir.parent)
                parts = relative_path.parts
                # Handle prototypes directory but exclude hybrid and kokkos
                if 'prototypes' in parts and 'hybrid' not in parts and 'kokkos' not in parts:
                    return True
        except ValueError:
            pass
        
        # Check file content for PETSc-specific keywords
        try:
            if config_path.suffix.lower() in ['.xml', '.json', '.yaml']:
                content = config_path.read_text().lower()
                petsc_keywords = ['petsc', 'ksp', 'mat', 'vec']
                if any(keyword in content for keyword in petsc_keywords):
                    # But exclude if it also contains hybrid/kokkos keywords
                    hybrid_keywords = ['kokkos', 'hybrid']
                    if not any(keyword in content for keyword in hybrid_keywords):
                        return True
        except Exception:
            pass
        
        return False
    
    def initialize(self, gsm_services: Optional['GSMServices'] = None) -> bool:
        """Initialize extension with GSM services."""
        try:
            # Check for PETSc installation
            petsc_dir = os.environ.get('PETSC_DIR')
            if not petsc_dir:
                logger.warning("PETSC_DIR not set in environment")
            
            # Validate that prototypes directory exists
            if not self.source_dir.exists():
                logger.error(f"Prototypes directory not found: {self.source_dir}")
                return False
            
            # Check if PETSc prototype files exist
            petsc_prototype = self.source_dir / "petsc_gcr_prototype.c"
            if not petsc_prototype.exists():
                logger.error(f"PETSc prototype not found: {petsc_prototype}")
                return False
            
            logger.info(f"PETSc-Only extension initialized with source dir: {self.source_dir}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize PETSc-Only extension: {e}")
            return False
    
    def cleanup(self) -> bool:
        """Cleanup extension resources."""
        # No specific cleanup needed
        return True